import numpy as np
from scipy.spatial.transform import Rotation as R
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

def generate_platform_points(diameter, num_points=6, offset_angle=0):
    """生成平台铰链点坐标"""
    radius = diameter / 2
    angles = np.linspace(0, 2*np.pi, num_points, endpoint=False) + offset_angle
    points = np.array([(radius*np.cos(theta), radius*np.sin(theta), 0) for theta in angles])
    return np.vstack([points, points[0]])

def calculate_workspace(lower_pts, upper_local, rotation, min_len, max_len, grid_size=30):
    """计算可达工作空间"""
    # 确定采样范围
    max_reach = max_len * 0.8  # 保守估计
    x = np.linspace(-max_reach, max_reach, grid_size)
    y = np.linspace(-max_reach, max_reach, grid_size)
    z = np.linspace(0, 2*max_reach, grid_size)
    X, Y, Z = np.meshgrid(x, y, z)
    
    # 向量化计算
    positions = np.vstack([X.ravel(), Y.ravel(), Z.ravel()]).T
    upper_global = (rotation @ upper_local[:-1].T).T + positions[:, np.newaxis, :]
    lower_expanded = lower_pts[np.newaxis, :, :]
    
    # 计算所有支腿长度
    distances = np.linalg.norm(upper_global - lower_expanded, axis=2)
    valid_mask = np.all((distances >= min_len) & (distances <= max_len), axis=1)
    
    return positions[valid_mask]

def plot_workspace(valid_points, lower_dia, upper_dia, pos, rotation):
    """优化后的工作空间可视化"""
    fig = plt.figure(figsize=(12, 8))
    ax = fig.add_subplot(111, projection='3d')
    
    # 调整参数：增大点尺寸和透明度
    ax.scatter(valid_points[:,0], valid_points[:,1], valid_points[:,2], 
               c='dodgerblue',    # 更亮的蓝色
               alpha=0.5,        # 透明度从0.08增加到0.25
               s=25,              # 点尺寸从2增大到15
               edgecolors='navy', # 添加深蓝色边缘
               linewidths=0.1,
               label='Workspace')
    
    # 平台结构绘制
    plot_platform(ax, lower_dia, upper_dia, pos, rotation)
    
    # 优化视角和光照
    ax.view_init(elev=25, azim=45)
    ax.set_box_aspect([1,1,1])
    ax.xaxis.pane.set_alpha(0.8)
    ax.yaxis.pane.set_alpha(0.8)
    ax.zaxis.pane.set_alpha(0.8)
    plt.title("Stewart Platform Workspace", fontsize=14)
    plt.show()



def plot_platform(ax, lower_dia, upper_dia, position, rotation):
    """平台结构绘制"""
    # 下平台
    theta = np.linspace(0, 2*np.pi, 100)
    lower_circle = lower_dia/2 * np.column_stack([np.cos(theta), np.sin(theta), np.zeros_like(theta)])
    ax.plot(lower_circle[:,0], lower_circle[:,1], lower_circle[:,2], 'r--', alpha=0.6)
    
    # 上平台
    upper_local = generate_platform_points(upper_dia, 6, np.pi/6)[:-1]
    upper_global = (rotation @ upper_local.T).T + position
    upper_circle = (rotation @ (upper_dia/2 * np.column_stack([np.cos(theta), np.sin(theta), np.zeros_like(theta)]).T)).T + position
    ax.plot(upper_circle[:,0], upper_circle[:,1], upper_circle[:,2], 'b-', alpha=0.8)

# 参数输入
lower_dia = 200.0  # 示例输入
upper_dia = 150.0
min_len = 100.0
max_len = 180.0
direction = np.array([0, 0, 1.0])  # 垂直方向
twist = np.radians(0.0)  # 无扭转

# 生成平台结构
lower_pts = generate_platform_points(lower_dia)[:-1]
upper_local = generate_platform_points(upper_dia, offset_angle=np.pi/6)

# 计算旋转矩阵
direction_norm = direction / np.linalg.norm(direction)
align_rot, _ = R.align_vectors([[0,0,1]], [direction_norm])
twist_rot = R.from_rotvec(twist * direction_norm)
rotation = (twist_rot * align_rot).as_matrix()

# 计算工作空间
valid_points = calculate_workspace(lower_pts, upper_local, rotation, min_len, max_len, grid_size=40)

# 可视化
plot_workspace(valid_points, lower_dia, upper_dia, np.array([0,0,200]), rotation)
